package internal.inference;

import inference.IType.ITypeEquation;
import inference.IType.ITypeEquationResolver;
import inference.IType.ITypeEquationVisitor;
import inference.IType.ITypeVariable;
import inference.IType.Type;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;

import parser.IParse.IArithmeticExpression;
import parser.IParse.IAssignment;
import parser.IParse.IBooleanExpression;
import parser.IParse.IDecimalLiteral;
import parser.IParse.IFunction;
import parser.IParse.IIntegerLiteral;
import parser.IParse.IParseNode;
import parser.IParse.IReturn;
import parser.IParse.IVariable;

class TypeEquationResolver implements ITypeEquationResolver {

	private static final boolean PRINT_ITERATIONS = true;
	private final List<ITypeEquation> equations;
	private final List<ITypeEquation> resolved;
	private final IFunction function;
	private final Map<String, Type> varToType;
	private final Map<ITypeVariable, IParseNode> varToNode;

	TypeEquationResolver(final IFunction function,
			final ITypeEquationVisitor visitor) {

		this.equations = visitor.getEquations();
		this.function = function;
		this.resolved = new ArrayList<ITypeEquation>();
		this.varToNode = visitor.getVariableMap();
		this.varToType = new HashMap<String, Type>();
	}

	public void printEquations() {

		for (final ITypeEquation eq : resolved) {
			final String lhs = varToString(eq.getLHS());
			final String exp1 = varToString(eq.getExpression1());
			if (eq.getExpression2() == null) {
				System.out.format("  %s = %s\n", lhs, exp1);
			} else {
				final String exp2 = varToString(eq.getExpression2());
				System.out.format("  %s = %s %s %s\n", lhs, exp1, eq
						.getOperator().toString(), exp2);
			}
		}
	}

	private String varToString(final ITypeVariable var) {

		return varType(var) != Type.UNDEFINED ? varType(var).toString() : var
				.toString();
	}

	private Type varType(final ITypeVariable variable) {
		return varToType.get(variable.getId());
	}

	@Override
	public void solveEquations() throws TypeException {
		ITypeVariable lhs = null;
		ITypeVariable exp1 = null;
		ITypeVariable exp2 = null;
		Type type1 = null;
		Type type2 = null;
		Iterator<ITypeEquation> eqnIter = null;
		ITypeEquation eqn = null;
		String op = null;
		IParseNode lhsNode = null;
		boolean isResolved = true;
		int varToTypeSize = varToType.size();
		int varToTypeSizeOld;
		int iternCount = 1;
		char nextPolyType = 'A';

		do {
			varToTypeSizeOld = varToTypeSize;
			eqnIter = equations.iterator();

			// A single iteration
			while (eqnIter.hasNext()) {
				eqn = eqnIter.next();
				lhs = eqn.getLHS();
				exp1 = eqn.getExpression1();
				exp2 = eqn.getExpression2();
				lhsNode = varToNode.get(lhs);
				op = eqn.getOperator();
				isResolved = false;

				// Exp1 Literal check
				if (varToNode.get(exp1) instanceof IIntegerLiteral) {
					varToType.put(exp1.getId(), Type.INTEGER);
				} else if (varToNode.get(exp1) instanceof IDecimalLiteral) {
					varToType.put(exp1.getId(), Type.REAL);
				}
				type1 = varToType.get(exp1.getId());

				if (exp2 != null) {
					// Exp2 Literal check
					if (varToNode.get(exp2) instanceof IIntegerLiteral) {
						varToType.put(exp2.getId(), Type.INTEGER);
					} else if (varToNode.get(exp2) instanceof IDecimalLiteral) {
						varToType.put(exp2.getId(), Type.REAL);
					}
					type2 = varToType.get(exp2.getId());
				}

				if (type1 != null
						&& (lhsNode instanceof IAssignment || lhsNode instanceof IReturn)) {
					// Assign or Return
					varToType.put(lhs.getId(), type1);
					isResolved = true;
				} else if (lhsNode instanceof IBooleanExpression) {
					// Boolean Op Check
					varToType.put(lhs.getId(), Type.BOOLEAN);
					if (type1 != null && type2 == null) {
						varToType.put(exp2.getId(), type1);
						isResolved = true;
					} else if (type2 != null && type1 == null) {
						varToType.put(exp1.getId(), type2);
						isResolved = true;
					} else if (type1 != type2) {
						throw new TypeException("Mismatch between types of "
								+ exp1.getId() + " and " + exp2.getId());
					}
				} else if (lhsNode instanceof IArithmeticExpression) {
					// Integer or Real
					if (type1 != null && type2 == null) {
						varToType.put(exp2.getId(), type1);
						varToType.put(lhs.getId(), type1);
						isResolved = true;
					} else if (type2 != null && type1 == null) {
						varToType.put(exp1.getId(), type2);
						varToType.put(lhs.getId(), type2);
						isResolved = true;
					} else if (type1 == type2 && type2 != null) {
						// Exp1 and exp2 already known
						varToType.put(lhs.getId(), type2);
						isResolved = true;
					} else if (type1 != type2) {
						throw new TypeException("Mismatch between types of "
								+ exp1.getId() + " and " + exp2.getId());
					}
				}
				// Remove resolved eqns
				if (isResolved) {
					resolved.add(eqn);
					eqnIter.remove();
				}
			}
			varToTypeSize = varToType.size();
			if (!resolved.isEmpty()) {
				System.out.println("Iteration #" + iternCount);
				printEquations();
			}
			iternCount++;

			if (varToTypeSizeOld == varToTypeSize) {
				// Polymorphic Types
				for (IVariable var : function.parameters()) {
					type1 = varToType.get(var.varId());
					if (type1 == null) {
						type1 = Type.valueOf("POLY_" + nextPolyType++);
						varToType.put(var.varId(), type1);
						varToTypeSize++;
						break;
					}
				}
			}
		} while (!equations.isEmpty() && varToTypeSizeOld != varToTypeSize);

		// Inferred type
		int i = 0;
		int size = function.parameters().size();
		System.out.print("Final Inferred Type:\n  ");
		for (IVariable var : function.parameters()) {
			type1 = varToType.get(var.varId());
			if (type1 == null) {
				type1 = Type.valueOf("POLY_" + nextPolyType++);
			}
			System.out.print(type1 + (++i < size ? " * " : " --> "));
		}
		type1 = varToType.get("T_0");
		if (type1 == null) {
			type1 = Type.valueOf("POLY_" + nextPolyType++);
		}
		System.out.println(type1);
	}
}
